import torch
from thop import profile
from basicsr.archs.dualCodeBook_arch import DualCodeBookNet
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

def main():
    """
    该脚本用于测量DualCodeBookNet模型的参数量和计算量（FLOPs）。
    """
    # 1. 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'使用设备: {device}')

    # 2. 定义模型
    # 这里的参数需要和你实际使用的模型参数保持一致
    model = DualCodeBookNet(
        in_channel=3,
        codebook_params=[[32, 1024, 512]],
        gt_resolution=256,
        LQ_stage=True,
        use_quantize=True,
        use_semantic_loss=False,
        use_residual=False,
        scale_factor=1,
        former=False,
        quantize_way='min',
    ).to(device)
    model.eval()

    # 3. 创建一个虚拟输入张量
    # 输入尺寸应与模型训练和推理时使用的尺寸一致
    # 这里我们假设输入图像的尺寸是 256x256
    # 尺寸为 [batch_size, channels, height, width]
    dummy_input = torch.randn(1, 3, 256, 256).to(device)

    # 4. 使用thop.profile计算参数量和FLOPs
    # 注意：这里的 FLOPs 通常指的是 MACs (Multiply-Accumulate operations)，
    # 有些地方会将 FLOPs 定义为 MACs 的两倍。thop 默认计算 MACs。
    macs, params = profile(model, inputs=(dummy_input, None), verbose=False)

    # 5. 打印结果
    # 将参数量转换为以“百万”（M）为单位
    params_m = params / 1e6
    # 将MACs转换为以“吉”（G）为单位
    macs_g = macs / 1e9

    print("\n模型计算量与参数量分析:")
    print("------------------------------------------")
    print(f"输入尺寸: {dummy_input.shape}")
    print(f"参数量 (Parameters): {params_m:.4f} M")
    print(f"计算量 (MACs): {macs_g:.4f} G")
    print("------------------------------------------")
    print("\n说明:")
    print(" - 参数量 (Parameters): 模型中所有可学习参数的总数，与输入尺寸无关。")
    print(" - 计算量 (MACs): 模型处理单个输入样本所需的乘法-累加操作次数，与输入尺寸相关。")

if __name__ == '__main__':
    main()